"""
EMG肌电伪迹去除器
使用小波包分解 + 频域特征识别方案去除肌电伪迹

作者：脑电波数据处理系统
版本：1.0
"""

import numpy as np
import logging
from collections import deque
from typing import Dict, List, Optional, Tuple
import pywt
from scipy import signal
from scipy.stats import entropy

# 配置日志
logger = logging.getLogger(__name__)

# 导入异常处理
from exception.base_exceptions import SignalProcessingException
from exception.exception_handler import global_exception_handler

class EMGArtifactRemover:
    """EMG肌电伪迹去除器
    
    使用小波包分解结合频域特征识别来检测和去除EMG伪迹
    EMG伪迹特征：
    - 频率范围：20-200Hz（主要在30-100Hz）
    - 高频能量占比高
    - 频谱熵较高
    - 幅度变化剧烈
    """
    
    def __init__(self, 
                 sampling_rate: int = 512,
                 window_size: int = 50,  # 降低窗口大小以提高敏感性
                 emg_threshold: float = 0.4,  # 降低阈值以提高灵敏度
                 wavelet: str = 'db4',
                 decomp_levels: int = 4):  # 调整分解层数以匹配窗口大小
        """初始化EMG伪迹去除器
        
        Args:
            sampling_rate: 采样率，默认512Hz
            window_size: 滑动窗口大小，默认100个样本点
            emg_threshold: EMG检测阈值，默认0.6
            wavelet: 小波基函数，默认'db4'
            decomp_levels: 小波包分解层数，默认6层
        """
        self.sampling_rate = sampling_rate
        self.window_size = window_size
        self.emg_threshold = emg_threshold
        self.wavelet = wavelet
        self.decomp_levels = decomp_levels
        
        # 数据缓冲区（8频域支持）
        self.buffer_size = max(window_size * 3, 300)  # 至少300个样本点
        self.raw_buffer = deque(maxlen=self.buffer_size)
        self.delta_buffer = deque(maxlen=self.buffer_size)
        self.theta_buffer = deque(maxlen=self.buffer_size)
        # 8频域缓冲区
        self.low_alpha_buffer = deque(maxlen=self.buffer_size)
        self.high_alpha_buffer = deque(maxlen=self.buffer_size)
        self.low_beta_buffer = deque(maxlen=self.buffer_size)
        self.high_beta_buffer = deque(maxlen=self.buffer_size)
        self.low_gamma_buffer = deque(maxlen=self.buffer_size)
        self.mid_gamma_buffer = deque(maxlen=self.buffer_size)
        # 向后兼容缓冲区
        self.alpha_buffer = deque(maxlen=self.buffer_size)
        self.beta_buffer = deque(maxlen=self.buffer_size)
        self.gamma_buffer = deque(maxlen=self.buffer_size)
        
        # 统计信息
        self.total_processed = 0
        self.emg_detected = 0
        self.artifacts_removed = 0
        
        # 频率范围定义
        self.freq_bands = {
            'delta': (0.5, 4),
            'theta': (4, 8),
            'alpha': (8, 13),
            'beta': (13, 30),
            'gamma': (30, 100),
            'emg': (20, 200)  # EMG主要频率范围
        }
        
        logger.info(f"EMG伪迹去除器初始化完成 - 采样率:{sampling_rate}Hz, 窗口:{window_size}, 阈值:{emg_threshold}")
    
    def process_signal(self, brain_wave_result) -> 'BrainWaveResult':
        """处理脑电信号，去除EMG伪迹
        
        Args:
            brain_wave_result: BrainWaveResult对象或字典
            
        Returns:
            处理后的BrainWaveResult对象
        """
        try:
            # 类型验证和自动转换
            if isinstance(brain_wave_result, dict):
                from response.BrainWaveResult import BrainWaveResult
                brain_wave_result = BrainWaveResult.from_dict(brain_wave_result)
                logger.debug("EMGArtifactRemover: 自动将字典转换为BrainWaveResult对象")
            elif not hasattr(brain_wave_result, 'delta'):
                raise SignalProcessingException(
                    f"EMGArtifactRemover: 输入数据类型无效: {type(brain_wave_result)}",
                    processing_step="emg_input_validation"
                )
            
            self.total_processed += 1
            
            # 更新数据缓冲区
            self._update_buffers(brain_wave_result)
            
            # 检查是否有足够的数据进行处理（小波包分解至少需要32个数据点）
            min_required_data = 32
            if len(self.raw_buffer) < min_required_data:
                # 数据不足，直接返回原始结果
                brain_wave_result.add_processing_step('EMGArtifactRemover', {
                    'status': 'insufficient_data',
                    'buffer_size': len(self.raw_buffer),
                    'required_size': min_required_data
                })
                return brain_wave_result
            
            # EMG伪迹检测
            emg_score = self._detect_emg_artifacts()
            
            if emg_score > self.emg_threshold:
                # 检测到EMG伪迹，进行去除
                self.emg_detected += 1
                processed_result = self._remove_emg_artifacts(brain_wave_result, emg_score)
                self.artifacts_removed += 1
                
                logger.debug(f"EMG伪迹已去除 - 评分:{emg_score:.3f}, 阈值:{self.emg_threshold}")
                return processed_result
            else:
                # 未检测到EMG伪迹，返回原始结果
                brain_wave_result.add_processing_step('EMGArtifactRemover', {
                    'status': 'no_emg_detected',
                    'emg_score': emg_score,
                    'threshold': self.emg_threshold
                })
                return brain_wave_result
                
        except Exception as e:
            processing_error = SignalProcessingException(
                f"EMG处理失败: {e}",
                processing_step="emg_artifact_removal"
            )
            global_exception_handler.handle_exception(processing_error)
            return brain_wave_result
    
    def _update_buffers(self, brain_wave_result):
        """更新数据缓冲区（8频域支持）"""
        self.raw_buffer.append(brain_wave_result.rawValue)
        self.delta_buffer.append(brain_wave_result.delta)
        self.theta_buffer.append(brain_wave_result.theta)
        
        # 8频域数据
        self.low_alpha_buffer.append(brain_wave_result.low_alpha)
        self.high_alpha_buffer.append(brain_wave_result.high_alpha)
        self.low_beta_buffer.append(brain_wave_result.low_beta)
        self.high_beta_buffer.append(brain_wave_result.high_beta)
        self.low_gamma_buffer.append(brain_wave_result.low_gamma)
        self.mid_gamma_buffer.append(brain_wave_result.mid_gamma)
        
        # 向后兼容数据
        self.alpha_buffer.append(brain_wave_result.alpha)
        self.beta_buffer.append(brain_wave_result.beta)
        self.gamma_buffer.append(brain_wave_result.gamma)
    
    def _detect_emg_artifacts(self) -> float:
        """检测EMG伪迹
        
        Returns:
            EMG评分 (0-1之间，越高表示EMG伪迹可能性越大)
        """
        try:
            # 获取最近的窗口数据
            raw_window = np.array(list(self.raw_buffer)[-self.window_size:])
            
            # 1. 小波包分解
            wp = pywt.WaveletPacket(raw_window, self.wavelet, maxlevel=self.decomp_levels)
            
            # 2. 计算各频段能量
            freq_energies = self._calculate_frequency_energies(wp, raw_window)
            
            # 3. 计算EMG特征
            emg_features = self._calculate_emg_features(raw_window, freq_energies)
            
            # 4. 综合评分
            emg_score = self._calculate_emg_score(emg_features)
            
            return min(max(emg_score, 0.0), 1.0)  # 限制在0-1范围内
            
        except Exception as e:
            detection_error = SignalProcessingException(
                f"EMG检测失败: {e}",
                processing_step="emg_detection"
            )
            global_exception_handler.handle_exception(detection_error)
            return 0.0
    
    def _calculate_frequency_energies(self, wp, raw_signal: np.ndarray) -> Dict[str, float]:
        """计算各频段能量"""
        try:
            # 计算频率分辨率
            freq_resolution = self.sampling_rate / len(raw_signal)
            
            # 获取所有叶子节点
            leaf_nodes = [node for node in wp.get_level(self.decomp_levels, 'natural')]
            
            energies = {}
            total_energy = 0
            
            for i, node in enumerate(leaf_nodes):
                # 计算该节点对应的频率范围
                freq_start = i * freq_resolution * (2 ** self.decomp_levels) / len(leaf_nodes)
                freq_end = (i + 1) * freq_resolution * (2 ** self.decomp_levels) / len(leaf_nodes)
                
                # 计算能量
                energy = np.sum(node.data ** 2)
                total_energy += energy
                
                # 分配到对应频段
                for band_name, (f_low, f_high) in self.freq_bands.items():
                    if freq_start >= f_low and freq_end <= f_high:
                        if band_name not in energies:
                            energies[band_name] = 0
                        energies[band_name] += energy
            
            # 归一化能量
            if total_energy > 0:
                for band in energies:
                    energies[band] /= total_energy
            
            return energies
            
        except Exception as e:
            energy_error = SignalProcessingException(
                f"频率能量计算失败: {e}",
                processing_step="frequency_energy_calculation"
            )
            global_exception_handler.handle_exception(energy_error)
            return {}
    
    def _calculate_emg_features(self, raw_signal: np.ndarray, freq_energies: Dict[str, float]) -> Dict[str, float]:
        """计算EMG特征"""
        features = {}
        
        try:
            # 1. 高频能量比 (EMG主要在高频)
            high_freq_energy = freq_energies.get('beta', 0) + freq_energies.get('gamma', 0) + freq_energies.get('emg', 0)
            total_energy = sum(freq_energies.values()) if freq_energies else 1
            features['high_freq_ratio'] = high_freq_energy / total_energy if total_energy > 0 else 0
            
            # 2. 频谱熵 (EMG具有较高的频谱熵)
            if freq_energies:
                energy_values = np.array(list(freq_energies.values()))
                energy_values = energy_values[energy_values > 0]  # 去除零值
                if len(energy_values) > 1:
                    features['spectral_entropy'] = entropy(energy_values)
                else:
                    features['spectral_entropy'] = 0
            else:
                features['spectral_entropy'] = 0
            
            # 3. 幅度变异系数 (EMG幅度变化剧烈)
            if len(raw_signal) > 1:
                features['amplitude_cv'] = np.std(raw_signal) / (np.mean(np.abs(raw_signal)) + 1e-8)
            else:
                features['amplitude_cv'] = 0
            
            # 4. 零交叉率 (EMG具有较高的零交叉率)
            zero_crossings = np.sum(np.diff(np.sign(raw_signal)) != 0)
            features['zero_crossing_rate'] = zero_crossings / len(raw_signal) if len(raw_signal) > 1 else 0
            
            # 5. 峰值因子 (EMG具有较高的峰值因子)
            rms = np.sqrt(np.mean(raw_signal ** 2))
            peak = np.max(np.abs(raw_signal))
            features['crest_factor'] = peak / (rms + 1e-8)
            
        except Exception as e:
            feature_error = SignalProcessingException(
                f"EMG特征计算失败: {e}",
                processing_step="emg_feature_calculation"
            )
            global_exception_handler.handle_exception(feature_error)
            features = {
                'high_freq_ratio': 0,
                'spectral_entropy': 0,
                'amplitude_cv': 0,
                'zero_crossing_rate': 0,
                'crest_factor': 0
            }
        
        return features
    
    def _calculate_emg_score(self, features: Dict[str, float]) -> float:
        """计算EMG综合评分"""
        try:
            # 特征权重 (根据EMG特征重要性设定)
            weights = {
                'high_freq_ratio': 0.3,      # 高频能量比 - 最重要
                'spectral_entropy': 0.25,    # 频谱熵
                'amplitude_cv': 0.2,         # 幅度变异系数
                'zero_crossing_rate': 0.15,  # 零交叉率
                'crest_factor': 0.1          # 峰值因子
            }
            
            # 特征归一化和评分
            score = 0.0
            
            # 1. 高频能量比评分 (>0.6认为是EMG)
            high_freq_score = min(features.get('high_freq_ratio', 0) / 0.6, 1.0)
            score += weights['high_freq_ratio'] * high_freq_score
            
            # 2. 频谱熵评分 (>2.0认为是EMG)
            entropy_score = min(features.get('spectral_entropy', 0) / 2.0, 1.0)
            score += weights['spectral_entropy'] * entropy_score
            
            # 3. 幅度变异系数评分 (>1.5认为是EMG)
            cv_score = min(features.get('amplitude_cv', 0) / 1.5, 1.0)
            score += weights['amplitude_cv'] * cv_score
            
            # 4. 零交叉率评分 (>0.1认为是EMG)
            zcr_score = min(features.get('zero_crossing_rate', 0) / 0.1, 1.0)
            score += weights['zero_crossing_rate'] * zcr_score
            
            # 5. 峰值因子评分 (>5.0认为是EMG)
            crest_score = min(features.get('crest_factor', 0) / 5.0, 1.0)
            score += weights['crest_factor'] * crest_score
            
            return score
            
        except Exception as e:
            score_error = SignalProcessingException(
                f"EMG评分计算失败: {e}",
                processing_step="emg_score_calculation"
            )
            global_exception_handler.handle_exception(score_error)
            return 0.0
    
    def _remove_emg_artifacts(self, brain_wave_result, emg_score: float) -> 'BrainWaveResult':
        """去除EMG伪迹"""
        try:
            # 创建结果副本
            result = brain_wave_result.copy()
            
            # 获取窗口数据进行小波去噪
            raw_window = np.array(list(self.raw_buffer)[-self.window_size:])
            
            # 小波去噪
            denoised_signal = self._wavelet_denoising(raw_window)
            
            # 计算去噪后的当前值 (取最后一个值)
            current_denoised = denoised_signal[-1]
            
            # 计算各频段的修正系数
            correction_factor = self._calculate_correction_factor(emg_score)
            
            # 记录原始值（8频域支持）
            original_values = {
                'rawValue': result.rawValue,
                'delta': result.delta,
                'theta': result.theta,
                'low_alpha': result.low_alpha,
                'high_alpha': result.high_alpha,
                'low_beta': result.low_beta,
                'high_beta': result.high_beta,
                'low_gamma': result.low_gamma,
                'mid_gamma': result.mid_gamma,
                # 向后兼容
                'alpha': result.alpha,
                'beta': result.beta,
                'gamma': result.gamma
            }
            
            # 应用EMG去除
            # 原始值使用去噪后的值
            result.rawValue = int(current_denoised)
            
            # 8频域根据EMG影响程度进行修正
            # Delta和Theta受影响较小
            result.delta = int(result.delta * (1 - correction_factor * 0.2))
            result.theta = int(result.theta * (1 - correction_factor * 0.3))
            
            # Alpha频段受中等影响
            result.low_alpha = int(result.low_alpha * (1 - correction_factor * 0.35))
            result.high_alpha = int(result.high_alpha * (1 - correction_factor * 0.45))
            
            # Beta频段受较大影响
            result.low_beta = int(result.low_beta * (1 - correction_factor * 0.7))
            result.high_beta = int(result.high_beta * (1 - correction_factor * 0.85))
            
            # Gamma频段受最大影响
            result.low_gamma = int(result.low_gamma * (1 - correction_factor * 0.85))
            result.mid_gamma = int(result.mid_gamma * (1 - correction_factor * 0.95))
            
            # 重新计算向后兼容的合成频段
            result.alpha = result.low_alpha + result.high_alpha
            result.beta = result.low_beta + result.high_beta
            result.gamma = result.low_gamma + result.mid_gamma
            
            # 重新计算分析指标
            result.update_analysis_indicators()
            
            # 添加处理步骤信息
            result.add_processing_step('EMGArtifactRemover', {
                'status': 'emg_removed',
                'emg_score': emg_score,
                'correction_factor': correction_factor,
                'original_values': original_values,
                'processed_values': {
                    'rawValue': result.rawValue,
                    'delta': result.delta,
                    'theta': result.theta,
                    'alpha': result.alpha,
                    'beta': result.beta,
                    'gamma': result.gamma
                }
            })
            
            return result
            
        except Exception as e:
            removal_error = SignalProcessingException(
                f"EMG伪迹去除失败: {e}",
                processing_step="emg_artifact_removal_process"
            )
            global_exception_handler.handle_exception(removal_error)
            return brain_wave_result
    
    def _wavelet_denoising(self, signal_data: np.ndarray) -> np.ndarray:
        """小波去噪"""
        try:
            # 小波分解
            coeffs = pywt.wavedec(signal_data, self.wavelet, level=self.decomp_levels)
            
            # 计算阈值 (使用软阈值)
            sigma = np.median(np.abs(coeffs[-1])) / 0.6745  # 噪声标准差估计
            threshold = sigma * np.sqrt(2 * np.log(len(signal_data)))
            
            # 应用软阈值
            coeffs_thresh = list(coeffs)
            coeffs_thresh[1:] = [pywt.threshold(detail, threshold, mode='soft') 
                                for detail in coeffs_thresh[1:]]
            
            # 重构信号
            denoised = pywt.waverec(coeffs_thresh, self.wavelet)
            
            # 确保长度一致
            if len(denoised) != len(signal_data):
                denoised = denoised[:len(signal_data)]
            
            return denoised
            
        except Exception as e:
            denoising_error = SignalProcessingException(
                f"小波去噪失败: {e}",
                processing_step="wavelet_denoising"
            )
            global_exception_handler.handle_exception(denoising_error)
            return signal_data
    
    def _calculate_correction_factor(self, emg_score: float) -> float:
        """计算修正系数"""
        # 根据EMG评分计算修正强度
        # 评分越高，修正越强
        return min(emg_score * 1.2, 0.8)  # 最大修正80%
    
    def get_statistics(self) -> Dict[str, float]:
        """获取统计信息"""
        detection_rate = self.emg_detected / self.total_processed if self.total_processed > 0 else 0
        removal_rate = self.artifacts_removed / self.emg_detected if self.emg_detected > 0 else 0
        
        return {
            'total_processed': self.total_processed,
            'emg_detected': self.emg_detected,
            'artifacts_removed': self.artifacts_removed,
            'detection_rate': detection_rate,
            'removal_rate': removal_rate,
            'current_threshold': self.emg_threshold
        }
    
    def reset_statistics(self):
        """重置统计信息"""
        self.total_processed = 0
        self.emg_detected = 0
        self.artifacts_removed = 0
        logger.info("EMG统计信息已重置")
    
    def update_threshold(self, new_threshold: float):
        """更新EMG检测阈值"""
        self.emg_threshold = max(0.1, min(new_threshold, 1.0))  # 限制在0.1-1.0范围内
        logger.info(f"EMG检测阈值已更新为: {self.emg_threshold}")

    def process_8bands(self, brain_wave_result) -> dict:
        """
        处理8频域脑波数据的EMG伪迹去除

        Args:
            brain_wave_result: 包含8频域数据的BrainWaveResult对象或字典

        Returns:
            dict: 经过EMG伪迹去除处理的结果字典
        """
        try:
            # 类型验证和转换
            if hasattr(brain_wave_result, 'to_dict'):
                # BrainWaveResult对象，转换为字典
                input_data = brain_wave_result.to_dict()
            elif isinstance(brain_wave_result, dict):
                # 已经是字典
                input_data = brain_wave_result.copy()
            else:
                logger.error(f"EMGArtifactRemover.process_8bands: 不支持的输入类型: {type(brain_wave_result)}")
                return {}
            
            # 对于EMG伪迹去除，我们使用现有的process_signal方法处理
            # 首先将字典转换回BrainWaveResult对象进行处理
            from response.BrainWaveResult import BrainWaveResult
            temp_result = BrainWaveResult.from_dict(input_data)
            
            # 使用现有的process_signal方法处理
            processed_result = self.process_signal(temp_result)
            
            # 将处理结果转换回字典格式
            result_dict = processed_result.to_dict()
            
            # 添加8频域处理信息
            if 'processing_info' not in result_dict:
                result_dict['processing_info'] = {}
            result_dict['processing_info']['emg_artifact_remover_8bands'] = {
                'status': 'processed',
                'details': 'EMG artifact removal applied using existing signal processing method.'
            }
            
            return result_dict
            
        except Exception as e:
            logger.error(f"EMG 8频域处理失败: {e}")
            # 返回原始数据字典
            if hasattr(brain_wave_result, 'to_dict'):
                return brain_wave_result.to_dict()
            elif isinstance(brain_wave_result, dict):
                return brain_wave_result.copy()
            else:
                return {}


# 便捷函数
def create_emg_remover(sampling_rate: int = 512, 
                      emg_threshold: float = 0.6) -> EMGArtifactRemover:
    """创建EMG伪迹去除器的便捷函数"""
    return EMGArtifactRemover(
        sampling_rate=sampling_rate,
        emg_threshold=emg_threshold
    )

def quick_emg_removal(brain_wave_result, 
                     emg_threshold: float = 0.6) -> 'BrainWaveResult':
    """快速EMG伪迹去除的便捷函数"""
    remover = create_emg_remover(emg_threshold=emg_threshold)
    return remover.process_signal(brain_wave_result)